package com.xwc.cloud.support.loadbalancer.core;


import com.xwc.cloud.support.loadbalancer.LoadbalancerSupportProperties;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.*;
import org.springframework.cloud.loadbalancer.core.*;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
import reactor.core.publisher.Mono;

import java.util.List;
import java.util.Random;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;


/**
 * 类描述：用于gateway的路由配置
 * 作者：徐卫超 (cc)
 * 时间 2022/5/12 17:25
 */

public class IpRoundRobinLoadBalancer implements ReactorServiceInstanceLoadBalancer {

    private static final Log log = LogFactory.getLog(RoundRobinLoadBalancer.class);
    final AtomicInteger position;
    final String serviceId;
    ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider;
    private LoadbalancerSupportProperties loadbalancerSupportProperties;

    public IpRoundRobinLoadBalancer(ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider,
                                    String serviceId, LoadbalancerSupportProperties loadbalancerSupportProperties) {
        this(serviceInstanceListSupplierProvider, serviceId, (new Random()).nextInt(1000), loadbalancerSupportProperties);
    }

    public IpRoundRobinLoadBalancer(ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider,
                                    String serviceId, int seedPosition, LoadbalancerSupportProperties loadbalancerSupportProperties) {
        this.serviceId = serviceId;
        this.serviceInstanceListSupplierProvider = serviceInstanceListSupplierProvider;
        this.position = new AtomicInteger(seedPosition);
        this.loadbalancerSupportProperties = loadbalancerSupportProperties;
    }

    public Mono<Response<ServiceInstance>> choose(Request request) {
        ServiceInstanceListSupplier supplier = (ServiceInstanceListSupplier) this.serviceInstanceListSupplierProvider.getIfAvailable(NoopServiceInstanceListSupplier::new);
        return supplier.get(request).next().map((serviceInstances) -> {
            return this.processInstanceResponse(supplier, serviceInstances, request);
        });
    }

    private Response<ServiceInstance> processInstanceResponse(ServiceInstanceListSupplier supplier,
                                                              List<ServiceInstance> serviceInstances, Request request) {
        Response<ServiceInstance> serviceInstanceResponse = this.getInstanceResponse(serviceInstances, request);
        if (supplier instanceof SelectedInstanceCallback && serviceInstanceResponse.hasServer()) {
            ((SelectedInstanceCallback) supplier).selectedServiceInstance((ServiceInstance) serviceInstanceResponse.getServer());
        }

        return serviceInstanceResponse;
    }

    private Response<ServiceInstance> getInstanceResponse(List<ServiceInstance> instances, Request request) {
        if (instances.isEmpty()) {
            if (log.isWarnEnabled()) {
                log.warn("No servers available for service: " + this.serviceId);
            }

            return new EmptyResponse();
        } else {
            LoadbalancerSupportProperties.IpRuleProperties ipRule = loadbalancerSupportProperties.getIpRule();
            String ip = getIp(request);
            ServiceInstance instance = null;
            // 获取指定路由信息
            if (StringUtils.hasText(ip)) {
                instance = instances.stream().filter(ins -> ins.getHost().equals(ip)).findFirst().orElse(null);
            }
            // 查看是否配置有默认路由
            if (instance == null && StringUtils.hasText(ipRule.getDefaultIp())) {
                instance = instances.stream().filter(ins -> ins.getHost().equals(ipRule.getDefaultIp())).findFirst()
                        .orElse(null);
            }
            if (instance == null) {
                // 检查是否配置了排除的路由节点
                if (loadbalancerSupportProperties.getIpRule().getExcludeIpRegex() == null) {
                    int pos = Math.abs(this.position.incrementAndGet());
                    instance = instances.get(pos % instances.size());
                } else {
                    // 获取排除外可用的路由节点
                    List<ServiceInstance> availableInstances = instances.stream()
                            .filter(ins -> !ipRule.getExcludeIpRegex().matcher(ins.getHost()).find())
                            .collect(Collectors.toList());
                    //  排除后没有路由节点可用 直接返回空的实例
                    if (availableInstances.isEmpty()) {
                        if (log.isWarnEnabled()) {
                            log.warn("No servers available for service: " + this.serviceId);
                        }
                        return new EmptyResponse();
                    }
                    int pos = Math.abs(this.position.incrementAndGet());
                    instance = availableInstances.get(pos % availableInstances.size());
                }
            }

            return new DefaultResponse(instance);
        }
    }

    public String getIp(Request<?> request) {
        Object context = request.getContext();
        if (context instanceof RequestDataContext) {
            RequestDataContext dataContext = (RequestDataContext) context;
            List<String> headerValues =
                    dataContext.getClientRequest().getHeaders().get(loadbalancerSupportProperties.getIpRule().getRouterAttrName());
            if (!CollectionUtils.isEmpty(headerValues)) {
                return headerValues.stream().findFirst().orElse(null);
            }
        }
        return null;
    }

}

